from typing import Dict, Any, TypedDict, Annotated
from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage, AIMessage

class ConditionalState(TypedDict):
    messages: Annotated[list, "对话消息列表"]
    result: Annotated[str, "处理结果"]
    decision: Annotated[str, "决策结果"]
    path: Annotated[str, "执行路径"]

class ConditionalGraphDemo:
    """条件分支图演示"""
    
    def __init__(self):
        self.graph = self._build_graph()
    
    def _build_graph(self):
        """构建条件分支图"""
        
        workflow = StateGraph(ConditionalState)
        
        # 添加节点
        workflow.add_node("analyze_input", self._analyze_input)
        workflow.add_node("make_decision", self._make_decision)
        workflow.add_node("handle_question", self._handle_question)
        workflow.add_node("handle_task", self._handle_task)
        workflow.add_node("handle_chat", self._handle_chat)
        workflow.add_node("format_output", self._format_output)
        
        # 设置入口点
        workflow.set_entry_point("analyze_input")
        
        # 设置条件边
        workflow.add_conditional_edges(
            "make_decision",
            self._route_decision,
            {
                "question": "handle_question",
                "task": "handle_task", 
                "chat": "handle_chat"
            }
        )
        
        # 设置其他边
        workflow.add_edge("analyze_input", "make_decision")
        workflow.add_edge("handle_question", "format_output")
        workflow.add_edge("handle_task", "format_output")
        workflow.add_edge("handle_chat", "format_output")
        workflow.add_edge("format_output", END)
        
        return workflow.compile()
    
    def _analyze_input(self, state: ConditionalState) -> ConditionalState:
        """分析输入"""
        messages = state["messages"]
        if not messages:
            messages = [HumanMessage(content="请帮我解答一个问题")]
        
        content = messages[-1].content.lower()
        
        return {
            "messages": messages,
            "result": f"分析输入: {content}",
            "decision": "",
            "path": "analyzed"
        }
    
    def _make_decision(self, state: ConditionalState) -> ConditionalState:
        """做出决策"""
        content = state["result"]
        
        # 简单的决策逻辑
        if "问题" in content or "解答" in content or "?" in content:
            decision = "question"
        elif "任务" in content or "执行" in content or "做" in content:
            decision = "task"
        else:
            decision = "chat"
        
        return {
            "messages": state["messages"],
            "result": content,
            "decision": decision,
            "path": f"decided: {decision}"
        }
    
    def _route_decision(self, state: ConditionalState) -> str:
        """路由决策"""
        return state["decision"]
    
    def _handle_question(self, state: ConditionalState) -> ConditionalState:
        """处理问题"""
        return {
            "messages": state["messages"],
            "result": "这是问题处理路径的响应",
            "decision": state["decision"],
            "path": state["path"] + " -> question_handled"
        }
    
    def _handle_task(self, state: ConditionalState) -> ConditionalState:
        """处理任务"""
        return {
            "messages": state["messages"],
            "result": "这是任务处理路径的响应",
            "decision": state["decision"],
            "path": state["path"] + " -> task_handled"
        }
    
    def _handle_chat(self, state: ConditionalState) -> ConditionalState:
        """处理聊天"""
        return {
            "messages": state["messages"],
            "result": "这是聊天处理路径的响应",
            "decision": state["decision"],
            "path": state["path"] + " -> chat_handled"
        }
    
    def _format_output(self, state: ConditionalState) -> ConditionalState:
        """格式化输出"""
        return {
            "messages": state["messages"],
            "result": f"最终结果: {state['result']} (路径: {state['path']})",
            "decision": state["decision"],
            "path": state["path"]
        }
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行条件分支图"""
        initial_state = {
            "messages": [HumanMessage(content=input_data.get("message", "请帮我解答一个问题"))],
            "result": "",
            "decision": "",
            "path": ""
        }
        
        result = self.graph.invoke(initial_state)
        
        return {
            "final_state": result,
            "decision": result["decision"],
            "execution_path": result["path"],
            "final_result": result["result"]
        }
